use crate::consts::{
    MAX_ALIGN_SIZE, PAGE_HUGE_ALIGN, SEGMENT_MASK, SEGMENT_SHIFT, SEGMENT_SIZE,
    SMALL_PAGES_PER_SEGMENT,
};
use crate::page::{Page, PageKind};
use crate::utils::access_guard::AccessGuard;
use crate::utils::double_linked_list::DoubleLinkedList;
use crate::utils::os_box::{OsBox, OsBoxBody};
use crate::utils::thread_id;
use std::alloc::Layout;
use std::cell::RefCell;
use std::mem;
use std::ptr::{addr_of, NonNull};
use std::thread::ThreadId;

pub(crate) struct Segment {
    // 记录在这个 segment 中使用了多少页。
    // 内部很难维护这个值，打个洞让 Heap 去维护他。
    pub(crate) page_used: usize,

    bin: usize,
    length: usize,
    capacity: usize,
    segment_size: usize,
    page_shift: usize,
    page_size: usize,
    thread_id: ThreadId,
    page_kind: PageKind,
    next: Option<SegmentOwner>,
    prev: Option<NonNull<Segment>>,
}

impl OsBoxBody for Segment {
    fn init() -> Self {
        Self {
            page_used: 0,
            bin: 0,
            length: 0,
            capacity: 0,
            segment_size: 0,
            page_shift: 0,
            page_size: 0,
            thread_id: thread_id(),
            page_kind: PageKind::Small,
            next: None,
            prev: None,
        }
    }

    fn before_drop(&mut self) {
        debug_assert_eq!(self.page_used, 0);
        debug_assert!(self.next.is_none());
        debug_assert!(self.prev.is_none());

        #[cfg(debug_assertions)]
        {
            use crate::utils::double_linked_list::DoubleLinkedNode;
            use crate::utils::single_linked_list::SingleLinkedNode;
            for i in 0..self.length {
                let page_ptr = self.page_at(i);
                let page = self.access(page_ptr);
                assert!(page.is_empty());
                assert!(page.prev().is_none());
                assert!(DoubleLinkedNode::next(page).is_none());
                assert!(SingleLinkedNode::next(page).is_none());
            }
        }
    }
}

#[repr(transparent)]
pub(crate) struct SegmentOwner(pub(crate) OsBox<Segment>);

impl Segment {
    pub(crate) fn new(
        required: usize,
        kind: PageKind,
        bin: usize,
        page_shift: usize,
    ) -> Option<(SegmentOwner, DoubleLinkedList<Page>)> {
        let capacity = match kind {
            PageKind::Small | PageKind::Large => {
                debug_assert_eq!(required, 0);
                let page_size = 1 << page_shift;
                let capacity = SEGMENT_SIZE / page_size;
                debug_assert_eq!(SEGMENT_SIZE % page_size, 0);
                debug_assert!(capacity >= 1 && capacity <= SMALL_PAGES_PER_SEGMENT);
                capacity
            }
            PageKind::Huge => {
                debug_assert!(page_shift == SEGMENT_SHIFT && required > 0);
                1
            }
        };

        let min_size = size_of::<Segment>() + (capacity * size_of::<Page>()) + 16 /* padding */;
        let header_size = Layout::from_size_align(min_size, MAX_ALIGN_SIZE)
            .unwrap()
            .pad_to_align()
            .size();
        let segment_size = if required == 0 {
            SEGMENT_SIZE
        } else {
            Layout::from_size_align(required + header_size, PAGE_HUGE_ALIGN)
                .unwrap()
                .pad_to_align()
                .size()
        };
        debug_assert!(segment_size >= required);
        let page_size = match kind {
            PageKind::Small | PageKind::Large => 1 << page_shift,
            PageKind::Huge => segment_size,
        };

        // TODO try to get it from caches
        // Allocate the segment
        let layout = Layout::from_size_align(segment_size, SEGMENT_SIZE)
            .unwrap()
            .pad_to_align();
        let mut segment = match OsBox::<Segment>::new_with_layout(layout) {
            None => return None,
            Some(segment) => segment,
        };
        segment.segment_size = segment_size;
        segment.page_shift = page_shift;
        segment.page_kind = kind;
        segment.page_size = page_size;
        segment.capacity = capacity;
        segment.bin = bin;
        let free_pages = segment.extend_pages();
        Some((SegmentOwner(segment), free_pages))
    }

    pub(crate) fn capacity(&self) -> usize {
        self.capacity
    }

    pub(crate) fn page_at(&self, idx: usize) -> NonNull<Page> {
        debug_assert!(idx < self.capacity);
        let page_ptr = unsafe { NonNull::from(self).add(1).cast::<Page>().add(idx) };
        debug_assert_eq!(
            (page_ptr.as_ptr() as usize - addr_of!(*self) as usize - size_of::<Segment>())
                / size_of::<Page>(),
            idx
        );
        debug_assert_eq!(
            (page_ptr.as_ptr() as usize - addr_of!(*self) as usize - size_of::<Segment>())
                % size_of::<Page>(),
            0
        );
        page_ptr
    }

    fn page_area_at(&self, idx: usize) -> (NonNull<u8>, usize) {
        let mut page_size = self.page_size;
        let page_area = unsafe {
            let mut page_area = NonNull::from(self).cast::<u8>().add(idx * self.page_size);
            if idx == 0 {
                let meta_length = size_of::<Segment>() + size_of::<Page>() * self.capacity;
                page_area = page_area.add(meta_length);
                page_size -= meta_length;
            }
            page_area
        };
        debug_assert_eq!(self.ptr_page(page_area.as_ptr()), self.page_at(idx));
        debug_assert!(
            page_area.as_ptr() as usize
                >= NonNull::from(self).as_ptr() as usize
                    + size_of::<Segment>()
                    + size_of::<Page>() * self.capacity
        );
        debug_assert!(
            page_area.as_ptr() as usize + page_size
                <= NonNull::from(self).as_ptr() as usize + SEGMENT_SIZE
        );
        (page_area, page_size)
    }

    /// 扩展 page, 一次性扩展所有页
    fn extend_pages(&mut self) -> DoubleLinkedList<Page> {
        let free_pages = DoubleLinkedList::default();
        if self.length >= self.capacity {
            return free_pages;
        }
        while self.length < self.capacity {
            let idx = self.length;
            self.length += 1;
            let page_ptr = self.page_at(idx);
            let (page_area, page_size) = self.page_area_at(idx);
            Page::init(page_ptr, self.bin, page_size, page_area);
            free_pages.push(self, page_ptr);
        }
        free_pages
    }

    pub(crate) fn is_local(&self) -> bool {
        self.thread_id == thread_id()
    }

    pub(crate) fn is_empty(&self) -> bool {
        self.page_used == 0
    }
    pub(crate) fn ptr_page(&self, ptr: *mut u8) -> NonNull<Page> {
        debug_assert!(!ptr.is_null());
        debug_assert_eq!(
            ptr as usize & !SEGMENT_MASK,
            NonNull::from(self).as_ptr() as usize
        );
        // TODO huge pages not support
        // if (segment->page_size > MI_SEGMENT_SIZE) return &segment->pages[0];  // huge pages
        let diff = ptr as usize & SEGMENT_MASK;
        debug_assert!(diff < SEGMENT_SIZE);
        let idx = diff >> self.page_shift;
        debug_assert!(idx < self.length);
        debug_assert!(idx == 0 || self.page_kind == PageKind::Small);
        self.page_at(idx)
    }
}

#[repr(transparent)]
pub(crate) struct SegmentGroup {
    head: RefCell<Option<SegmentOwner>>,
}

impl Default for SegmentGroup {
    fn default() -> Self {
        Self {
            head: RefCell::new(None),
        }
    }
}

impl AccessGuard<Page> for Segment {
    fn before_access(&self, ptr: NonNull<Page>) {
        debug_assert_eq!(
            (ptr.as_ptr() as usize - NonNull::from(self).as_ptr() as usize - size_of::<Segment>())
                % size_of::<Page>(),
            0
        );
        debug_assert!(ptr.as_ptr() as usize >= self.page_at(0).as_ptr() as usize);
        debug_assert!(ptr.as_ptr() as usize <= self.page_at(self.length - 1).as_ptr() as usize);
    }
}

impl SegmentGroup {
    #[cfg(test)]
    pub(crate) fn len(&self) -> usize {
        let mut node_idx = 0;
        let mut current = self.head.borrow().as_ref().map(|v| v.0.as_ptr());
        while let Some(seg_node) = current {
            let node = unsafe { seg_node.as_ref() };
            node_idx += 1;
            current = node.next.as_ref().map(|v| v.0.as_ptr());
        }
        node_idx
    }
    pub(crate) fn push(&self, mut current: SegmentOwner) {
        debug_assert!(current.0.prev.is_none());
        debug_assert!(current.0.next.is_none());
        let current_ptr = current.0.as_ptr();
        let old_head = self.head.take();
        if let Some(mut old_head) = old_head {
            debug_assert!(old_head.0.prev.is_none());
            old_head.0.prev = Some(current_ptr);
            current.0.next = Some(old_head);
        }
        self.head.replace(Some(current));
    }

    pub(crate) fn leave<G>(&self, access_guard: &G, current_ptr: NonNull<Segment>) -> SegmentOwner
    where
        G: AccessGuard<Segment>,
    {
        let current = access_guard.access(current_ptr);
        let mut next_ptr = mem::replace(&mut current.next, None);
        let prev_ptr = mem::replace(&mut current.prev, None);

        if let Some(next_ptr) = next_ptr.as_mut() {
            debug_assert_eq!(next_ptr.0.prev, Some(current_ptr));
            next_ptr.0.prev = prev_ptr;
        }
        match prev_ptr {
            None => {
                let head = self.head.take();
                debug_assert!(head.is_some());
                let head = head.unwrap();
                debug_assert_eq!(head.0.as_ptr(), current_ptr);
                self.head.replace(next_ptr);
                head
            }
            Some(prev_ptr) => {
                let prev = access_guard.access(prev_ptr);
                let prev_next_ptr = prev.next.take();
                debug_assert!(prev_next_ptr.is_some());
                let prev_next_ptr = prev_next_ptr.unwrap();
                debug_assert_eq!(prev_next_ptr.0.as_ptr(), current_ptr);
                prev.next = next_ptr;
                prev_next_ptr
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::consts::SMALL_PAGE_SHIFT;
    use crate::page::PageKind;
    use crate::segment::{Segment, SegmentGroup};
    use crate::utils::access_guard::AccessGuard;
    use std::ptr::NonNull;

    fn assert_group<G>(
        guard: &G,
        group: &SegmentGroup,
        segs: &Vec<NonNull<Segment>>,
        expected: &[usize],
    ) where
        G: AccessGuard<Segment>,
    {
        let mut node_idx = 0;
        let mut current = group.head.borrow().as_ref().map(|v| v.0.as_ptr());
        while let Some(seg_node) = current {
            let node = guard.access(seg_node);
            assert_eq!(seg_node, segs[expected[node_idx]]);
            match node.prev {
                None => assert_eq!(group.head.borrow().as_ref().unwrap().0.as_ptr(), seg_node),
                Some(prev_ptr) => {
                    let prev_next_ptr = guard.access(prev_ptr).next.as_ref().unwrap().0.as_ptr();
                    assert_eq!(prev_next_ptr, seg_node);
                }
            }
            if let Some(next_ptr) = node.next.as_ref().map(|v| v.0.as_ptr()) {
                let next_prev_ptr = guard.access(next_ptr).prev.unwrap();
                assert_eq!(next_prev_ptr, seg_node);
            }
            node_idx += 1;
            current = node.next.as_ref().map(|v| v.0.as_ptr());
        }
        assert_eq!(node_idx, expected.len());
        assert_eq!(expected.is_empty(), group.head.borrow().is_none());
    }

    struct TestSegmentGuard;

    impl AccessGuard<Segment> for TestSegmentGuard {
        fn before_access(&self, _ptr: NonNull<Segment>) {}
    }
    #[test]
    fn test_segment_group() {
        let group = SegmentGroup::default();
        let mut segs = Vec::new();
        let guard = TestSegmentGuard;

        const TEST_BIN: usize = 24;
        (0..5).for_each(|_| {
            let (seg, pages) =
                Segment::new(0, PageKind::Small, TEST_BIN, SMALL_PAGE_SHIFT).unwrap();
            let seg_ptr = seg.0.as_ptr();
            pages.give_up_all(seg.0.as_ref());
            group.push(seg);
            segs.push(seg_ptr);
        });
        assert_group(&guard, &group, &segs, &[4, 3, 2, 1, 0]);

        let seg = group.leave(&guard, segs[2]);
        assert_group(&guard, &group, &segs, &[4, 3, 1, 0]);
        group.push(seg);
        assert_group(&guard, &group, &segs, &[2, 4, 3, 1, 0]);

        // 将 Segment 逐个从 group 中移除
        segs.iter().for_each(|&seg_ptr| {
            // 逐个放弃所有权，否则 Segment drop 检查会报错
            drop(group.leave(&guard, seg_ptr));
        });
        assert_group(&guard, &group, &segs, &[]);
    }
}
